from abc import ABC
from enum import Enum
from typing import Generic, Optional, TypeVar

from aicenter.models import AIModel


class ContentType(Enum):
    """
    内容类型
    """
    Text = "text"
    Image = "image"


class MessageRole(Enum):
    """
    消息来源角色

    即，发出该消息的人、AI 或 实体。
    """
    System = "system"
    User = "user"
    AI = "ai"


class AIModelRequest(ABC):
    """
    AI 模型请求类
    """
    def __init__(self, model: AIModel, params: Optional[dict]):
        self._model: AIModel = model
        self._params = dict(params) if params else {}

    @property
    def model(self) -> AIModel:
        """AI 模型"""
        return self._model

    @property
    def params(self) -> dict:
        """额外参数"""
        return dict(self._params)


class AITokenUsage:
    """
    AI 令牌使用量类
    """
    def __init__(self, prompt_usage: int, input_usage: int, output_usage: int, total_usage: Optional[int] = None):
        self._prompt_usage: int = prompt_usage
        self._input_usage: int = input_usage
        self._output_usage: int = output_usage
        self._total_usage: int = (
            total_usage
            if total_usage is not None else
            prompt_usage + input_usage + output_usage
        )

    @property
    def prompt_usage(self) -> int:
        """提示词令牌使用量"""
        return self._prompt_usage

    @property
    def input_usage(self) -> int:
        """输入令牌使用量"""
        return self._input_usage

    @property
    def output_usage(self) -> int:
        """输出令牌使用量"""
        return self._output_usage

    @property
    def total_usage(self) -> int:
        """总令牌使用量"""
        return self._total_usage


T = TypeVar('T', bound=AIModelRequest)


class AIModelResponse(ABC, Generic[T]):
    """
    AI 模型响应类
    """
    def __init__(self, request: T, token_usage: AITokenUsage):
        self._request = request
        self._token_usage = token_usage

    @property
    def model(self) -> AIModel:
        """AI 模型"""
        return self._request.model

    @property
    def params(self) -> dict:
        """额外参数"""
        return self._request.params

    @property
    def token_usage(self) -> AITokenUsage:
        """本次令牌使用量"""
        return self._token_usage
